#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18

"""
【节点名称】：
    ImgDirReaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    fps (int): 输出帧率
【备注】：
    无
"""
import os
import threading
import time
import cv2
import uuid
import os
import platform
import argparse
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, QoS, BaseNode, get_extra_args, Rate


class ImgDirReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True,
            **kwargs
        )
        self.image_dir = self.get_param("image_dir", "G:/dataset/BB_Object365/images/val")
        self.repeated = self.get_param("repeated", 0)
        self.auto_next = self.get_param("auto_next", False)
        self.rate = self.get_param("rate", 10)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.img_i = 0
        self.frame_id = 0
        self.done = False

        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        
        files_and_directories = os.listdir(self.image_dir)
        files_and_directories = sorted(files_and_directories)
        self.image_fns = []
        for img_name in files_and_directories:
            if os.path.splitext(img_name)[-1] in ['.jpg', '.jpeg', '.png']:
                self.image_fns.append(os.path.join(self.image_dir, img_name))
        # print(self.image_fns)
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.image_fns):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1
                if self.repeated and self.img_i == len(self.image_fns):
                    self.img_i = 0
            else:
                self.done = True

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img = cv2.imread(self.image_fns[img_i])

            if self.b_use_shm:
                msg = self._image_writer.cvimg2sms_mem(img)
            else:
                msg = cvimg2sms(img)

            msg['img_id'] = img_i
            msg['img_total'] = len(self.image_fns)
            msg['file_name'] = os.path.basename(self.image_fns[img_i])
            msg['client_id'] = self.client_id
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = ImgDirReaderNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.launch_next()
    if node.auto_next:
        r = Rate(node.rate)
        while not node.done:
            node.launch_next()
            r.sleep()
        node.shutdown()
    node.spin()
